import React, { useState } from 'react';
import { PageContainer } from '@ant-design/pro-components';
import { Card, Typography, Alert, App, Avatar, Switch, Space, Divider } from 'antd';
import { Bubble, Sender, XProvider, useXAgent, useXChat } from '@ant-design/x';
import { UserOutlined, RobotOutlined, CodeOutlined } from '@ant-design/icons';
import SyntaxHighlighter from 'react-syntax-highlighter';
import { docco } from 'react-syntax-highlighter/dist/esm/styles/hljs';

const { Title, Paragraph, Text } = Typography;

// 定义消息类型
interface XMessage {
  id: string;
  message: string;
  status: 'local' | 'remote' | 'requesting';
}

// 直接在代码中设置 API Key
// 注意：在实际生产环境中，应该使用环境变量或安全的配置管理系统来存储 API Key
// 这里仅作为示例，实际应用中不建议将 API Key 硬编码在代码中
const DEEPSEEK_API_KEY = 'sk-b0cdf2284e414d2793799ab2b4eafb5d'; // 请替换为您的实际 API Key

// 定义角色样式
const roles = {
  user: {
    placement: 'end',
    avatar: { icon: <UserOutlined />, style: { background: '#87d068' } },
  },
  assistant: {
    placement: 'start',
    avatar: { icon: <RobotOutlined />, style: { background: '#fde3cf' } },
  },
};

// 系统提示词，用于指导模型输出 JSON 格式
const JSON_SYSTEM_PROMPT = `
请以 JSON 格式回答用户的问题。
回答应该包含以下字段：
- "answer": 对用户问题的直接回答
- "explanation": 对回答的详细解释
- "references": 相关参考资料或来源（如有）

示例输出格式：
{
  "answer": "简短的回答",
  "explanation": "详细的解释",
  "references": ["参考资料1", "参考资料2"]
}
`;

const AntDesignXPage: React.FC = () => {
  const [error, setError] = useState<string>('');
  const [inputValue, setInputValue] = useState<string>('');
  const [useJsonOutput, setUseJsonOutput] = useState<boolean>(false);
  const { message } = App.useApp();

  // 自定义请求函数，处理 DeepSeek API 请求
  const customRequest = async ({ message: userMessage }: { message: string }, { onSuccess, onUpdate, onError }: any) => {
    console.log('Custom request called with message:', userMessage);
    
    try {
      // 获取当前对话历史
      const apiMessages = messages
        .filter(msg => msg.status !== 'requesting')
        .map(msg => ({
          role: msg.status === 'local' ? 'user' : 'assistant',
          content: msg.message,
        }));
      
      // 如果启用了 JSON 输出，添加系统消息
      if (useJsonOutput && !apiMessages.some(msg => msg.role === 'system')) {
        apiMessages.unshift({
          role: 'system',
          content: JSON_SYSTEM_PROMPT,
        });
      }
      
      // 添加当前用户消息
      apiMessages.push({
        role: 'user',
        content: userMessage,
      });
      
      console.log('Sending to DeepSeek API:', apiMessages);
      
      // 准备请求配置
      const requestBody: any = {
        model: 'deepseek-chat',
        messages: apiMessages,
        stream: true, // 启用流式响应
      };
      
      // 如果启用了 JSON 输出，添加 response_format 参数
      if (useJsonOutput) {
        requestBody.response_format = { type: 'json_object' };
        requestBody.max_tokens = 2048; // 设置较大的 max_tokens 防止 JSON 被截断
      }
      
      // 使用 fetch 调用 DeepSeek API
      const response = await fetch('https://api.deepseek.com/v1/chat/completions', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${DEEPSEEK_API_KEY}`,
        },
        body: JSON.stringify(requestBody),
      });
      
      if (!response.ok) {
        const errorText = await response.text();
        console.error('DeepSeek API error:', response.status, errorText);
        throw new Error(`API 返回错误: ${response.status} - ${errorText}`);
      }
      
      // 处理流式响应
      const reader = response.body?.getReader();
      if (!reader) {
        throw new Error('无法获取响应流');
      }
      
      const decoder = new TextDecoder();
      let fullContent = '';
      
      // 读取流式数据
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        
        const chunk = decoder.decode(value, { stream: true });
        console.log('Received chunk:', chunk);
        
        // 解析 SSE 数据
        const lines = chunk.split('\n\n');
        for (const line of lines) {
          if (line.startsWith('data: ') && !line.includes('[DONE]')) {
            try {
              const data = JSON.parse(line.substring(6));
              if (data.choices && data.choices[0] && data.choices[0].delta && data.choices[0].delta.content) {
                const content = data.choices[0].delta.content;
                fullContent += content;
                onUpdate(fullContent);
              }
            } catch (e) {
              console.warn('Error parsing SSE data:', e);
            }
          }
        }
      }
      
      // 完成后调用 onSuccess
      if (fullContent) {
        onSuccess(fullContent);
      } else {
        onSuccess('我无法生成回复，请稍后再试。');
      }
    } catch (err: any) {
      console.error('Error in custom request:', err);
      setError(`API 请求失败: ${err.message || '未知错误'}`);
      message.error('调用 DeepSeek API 时出错');
      
      // 调用 onError 回调
      onError(err);
      
      // 返回一个错误消息
      onSuccess(`出错了: ${err.message || '未知错误'}。请检查网络连接和 API Key。`);
    }
  };

  // 创建 AI 代理
  const [agent] = useXAgent({
    request: customRequest,
    onError: (err) => {
      console.error('Agent error:', err);
      setError(`API 调用失败: ${err.message || '未知错误'}`);
      message.error('调用 DeepSeek API 时出错');
    }
  });

  // 使用 useXChat 管理对话
  const { messages, onRequest } = useXChat<XMessage>({
    agent,
    initialMessages: [
      {
        id: '1',
        message: '你好，我是 AI 助手！我基于 DeepSeek 模型构建。请输入您的问题。',
        status: 'remote',
      }
    ],
  });

  // 处理消息发送
  const handleSubmit = (content: string) => {
    if (!content.trim()) return;
    
    try {
      // 调用 onRequest 发送请求
      console.log('Submitting message:', content);
      onRequest(content);
      setInputValue('');
    } catch (err: any) {
      console.error('Error in handleSubmit:', err);
      setError(`发送消息失败: ${err.message || '未知错误'}`);
      message.error('发送消息时出错');
    }
  };

  // 格式化 JSON 消息内容
  const formatMessageContent = (content: string, status: string) => {
    if (status !== 'remote' || !useJsonOutput) {
      return content;
    }
    
    try {
      // 尝试解析 JSON
      if (content.trim().startsWith('{') && content.trim().endsWith('}')) {
        const jsonData = JSON.parse(content);
        return (
          <div>
            {jsonData.answer && (
              <div>
                <Text strong>回答：</Text>
                <div>{jsonData.answer}</div>
                <Divider style={{ margin: '8px 0' }} />
              </div>
            )}
            
            {jsonData.explanation && (
              <div>
                <Text strong>解释：</Text>
                <div>{jsonData.explanation}</div>
                <Divider style={{ margin: '8px 0' }} />
              </div>
            )}
            
            {jsonData.references && jsonData.references.length > 0 && (
              <div>
                <Text strong>参考资料：</Text>
                <ul style={{ paddingLeft: 20, margin: '4px 0' }}>
                  {jsonData.references.map((ref: string, index: number) => (
                    <li key={index}>{ref}</li>
                  ))}
                </ul>
              </div>
            )}
            
            <Divider style={{ margin: '8px 0' }} />
            <div>
              <Text type="secondary" style={{ fontSize: '12px' }}>原始 JSON：</Text>
              <SyntaxHighlighter language="json" style={docco} customStyle={{ fontSize: '12px' }}>
                {JSON.stringify(jsonData, null, 2)}
              </SyntaxHighlighter>
            </div>
          </div>
        );
      }
    } catch (e) {
      console.warn('Error parsing JSON message:', e);
    }
    
    return content;
  };

  return (
    <PageContainer>
      <Card>
        <Typography>
          <Title level={2}>AI 助手演示 (Ant Design X + DeepSeek)</Title>
          <Paragraph>
            这是使用 Ant Design X 组件构建的 AI 助手演示，连接到 DeepSeek API，支持流式输出和 JSON 格式化输出。
          </Paragraph>
          <Space style={{ marginBottom: 16 }}>
            <Switch 
              checked={useJsonOutput} 
              onChange={setUseJsonOutput} 
              checkedChildren="JSON 输出" 
              unCheckedChildren="普通输出"
            />
            <Text type="secondary">
              {useJsonOutput ? '已启用 JSON 结构化输出' : '使用普通文本输出'}
            </Text>
          </Space>
          {error && (
            <Alert
              message="错误"
              description={error}
              type="error"
              showIcon
              style={{ marginBottom: 16 }}
              closable
              onClose={() => setError('')}
            />
          )}
          <div>
            <p>调试信息：</p>
            <p>消息数量: {messages.length}</p>
            <p>请求状态: {agent.isRequesting() ? '请求中' : '空闲'}</p>
            <p>API Key 状态: {DEEPSEEK_API_KEY ? '已设置' : '未设置'}</p>
            <p>输出模式: {useJsonOutput ? 'JSON 结构化输出' : '普通文本输出'}</p>
          </div>
        </Typography>
      </Card>

      <Card style={{ marginTop: 16, minHeight: 500 }}>
        <div style={{ height: 350, overflowY: 'auto', marginBottom: 16, padding: 8 }}>
          <Bubble.List
            roles={roles}
            items={messages.map(({ id, message, status }) => ({
              key: id,
              role: status === 'local' ? 'user' : 'assistant',
              content: formatMessageContent(message, status),
              loading: status === 'requesting',
            }))}
          />
        </div>
        
        <Sender
          value={inputValue}
          onChange={(v) => setInputValue(v)}
          onSubmit={handleSubmit}
          loading={agent.isRequesting()}
          placeholder={useJsonOutput ? "请输入问题，将以 JSON 格式回答..." : "请输入您的问题..."}
        />
      </Card>
    </PageContainer>
  );
};

// 使用 XProvider 和 App 包装导出的组件
export default () => (
  <App>
    <XProvider>
      <AntDesignXPage />
    </XProvider>
  </App>
);
